#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <math.h>
#include "gaussian.h"
#include "random.h"

static double gaussian_sample(double m, double s)
/* ========================================================================
 * Returns a normal (Gaussian) distributed real number.
 * NOTE: use s > 0.0
 *
 * Uses a very accurate approximation of the normal idf due to Odeh & Evans, 
 * J. Applied Statistics, 1974, vol 23, pp 96-97.
 * ========================================================================
 * copied from http://www.cs.wm.edu/~va/software/park/
 */
{ 
  const double p0 = 0.322232431088;     const double q0 = 0.099348462606;
  const double p1 = 1.0;                const double q1 = 0.588581570495;
  const double p2 = 0.342242088547;     const double q2 = 0.531103462366;
  const double p3 = 0.204231210245e-1;  const double q3 = 0.103537752850;
  const double p4 = 0.453642210148e-4;  const double q4 = 0.385607006340e-2;
  double u, t, p, q, z;

  u = RANDOM_DOUBLE;
  if (u < 0.5)
    t = sqrt(-2.0 * log(u));
  else
    t = sqrt(-2.0 * log(1.0 - u));
  p   = p0 + t * (p1 + t * (p2 + t * (p3 + t * p4)));
  q   = q0 + t * (q1 + t * (q2 + t * (q3 + t * q4)));
  if (u < 0.5)
    z = (p / q) - t;
  else
    z = t - (p / q);

  return (m + s * z);
}

static double gaussian_score(double mean, double var, double val)
{
  double diff;

  diff = val - mean;
  
  return - diff * diff / (2.0 * var) - .5 * log(2 * PI * var);
}


void UnivarGaussian_Init_Params(UnivarGaussian_t * dist, int nparams, ...)
{
  if (0 == nparams)
  {
    dist->fixed_mean = dist->fixed_var = 0;
  }
  else if (1 == nparams)
  {
    va_list ap;

    va_start(ap, nparams);

    dist->var = va_arg(ap, double);

    dist->fixed_mean = 0;
    dist->fixed_var = 1;
    
    va_end(ap);
  }  
  else if (2 == nparams)
  {
    va_list ap;

    va_start(ap, nparams);

    dist->mean = va_arg(ap, double);
    dist->var = va_arg(ap, double);

    dist->fixed_mean = dist->fixed_var = 1;
    
    va_end(ap);
  }
  else
  {
    fprintf(stderr, "error: UnivarGaussian takes at 0, 1, or 2 parameter\n");
    exit(1);
  }

  if (dist->fixed_var && (dist->var < 0))
  {
    fprintf(stderr, "error: UnivarGaussian: variance must be >= 0\n");
    exit(1);
  }
  
  if (dist->fixed_var)
    dist->sd = sqrt(dist->var);
}


double UnivarGaussian_Sample(const UnivarGaussian_t * dist, int nargs, ...)
{
  double mean;
  double sd;
  
  if (dist->fixed_mean && dist->fixed_var)
  {
    mean = dist->mean;
    sd = dist->sd;
  }
  else if (dist->fixed_var)
  {
    va_list ap;

    if (1 != nargs)
    {
      fprintf(stderr, "error: UnivarGaussian expecting only 1 parameter\n");
      exit(1);
    }

    va_start(ap, nargs);

    mean = va_arg(ap, double);

    va_end(ap);

    sd = dist->sd;
  }
  else
  {
    va_list ap;
    double var;

    if (2 != nargs)
    {
      fprintf(stderr, "error: UnivarGaussian expecting 2 parameters\n");
      exit(1);
    }

    va_start(ap, nargs);

    mean = va_arg(ap, double);
    var = va_arg(ap, double);

    va_end(ap);

    if (var < 0)
    {
      fprintf(stderr, "error: UnivarGaussian: variance must be >= 0\n");
      exit(1);
    }

    sd = sqrt(var);
  }
    
  return gaussian_sample(mean, sd);
}


static double UnivarGaussian_LogProb_Args(const UnivarGaussian_t * dist,
                                          double val, int nargs, va_list args)
{
  double mean;
  double var;
  
  if (dist->fixed_mean && dist->fixed_var)
  {
    mean = dist->mean;
    var = dist->var;
  }
  else if (dist->fixed_var)
  {
    if (1 != nargs)
    {
      fprintf(stderr, "error: UnivarGaussian expecting only 1 parameter\n");
      exit(1);
    }

    mean = va_arg(args, double);

    var = dist->var;
  }
  else
  {
    if (2 != nargs)
    {
      fprintf(stderr, "error: UnivarGaussian expecting 2 parameters\n");
      exit(1);
    }

    mean = va_arg(args, double);
    var = va_arg(args, double);

    if (var < 0)
    {
      fprintf(stderr, "error: UnivarGaussian: variance must be >= 0\n");
      exit(1);
    }
  }
    
  return gaussian_score(mean, var, val);
}

double UnivarGaussian_LogProb(const UnivarGaussian_t * dist, double val,
                              int nargs, ...)
{
  double retval;
  va_list args;

  va_start(args, nargs);
  retval = UnivarGaussian_LogProb_Args(dist, val, nargs, args);
  va_end(args);
  
  return retval;
}

double UnivarGaussian_Prob(const UnivarGaussian_t * dist, double val,
                           int nargs, ...)
{
  double retval;
  va_list args;

  va_start(args, nargs);
  retval = exp(UnivarGaussian_LogProb_Args(dist, val, nargs, args));
  va_end(args);
  
  return retval;
}

